package com.meiya.flink.util;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.net.URL;
import java.util.Collections;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * 类加载工具类，支持本地开发和集群环境下的类加载
 */
public class ClassLoaderUtil {
    private static final Logger logger = LoggerFactory.getLogger(ClassLoaderUtil.class);

    /**
     * 获取指定包下的所有类
     * 支持本地开发和集群环境
     *
     * @param packageName 包名
     *
     * @return 类集合
     */
    public static Set<Class<?>> getClasses(String packageName) {
        Set<Class<?>> classes = new HashSet<>();
        try {
            // 获取当前线程的类加载器
            ClassLoader classLoader = Thread.currentThread().getContextClassLoader();

            // 将包名转换为路径
            String packagePath = packageName.replace('.', '/');

            // 获取所有资源URL
            Enumeration<URL> resources = classLoader.getResources(packagePath);

            while (resources.hasMoreElements()) {
                URL resource = resources.nextElement();
                String protocol = resource.getProtocol();

                if ("file".equals(protocol)) {
                    // 本地开发环境
                    findClassesInDirectory(new File(resource.getFile()), packageName, classes);
                } else if ("jar".equals(protocol)) {
                    // 集群环境（JAR包）
                    findClassesInJar(resource, packageName, classes);
                }
            }

            logger.info("在包 {} 中找到 {} 个类", packageName, classes.size());
            return classes;

        } catch (Exception e) {
            logger.error("加载包 {} 中的类时发生错误: {}", packageName, e.getMessage(), e);
            return Collections.emptySet();
        }
    }

    /**
     * 在目录中查找类
     */
    private static void findClassesInDirectory(File directory, String packageName, Set<Class<?>> classes) {
        if (directory.exists() && directory.isDirectory()) {
            File[] files = directory.listFiles();
            if (files != null) {
                for (File file : files) {
                    if (file.isDirectory()) {
                        findClassesInDirectory(file, packageName + "." + file.getName(), classes);
                    } else if (file.getName().endsWith(".class")) {
                        try {
                            String className = packageName + "." + file.getName().substring(0, file.getName().length() - 6);
                            classes.add(Class.forName(className));
                        } catch (ClassNotFoundException e) {
                            logger.warn("无法加载类: {}", e.getMessage());
                        }
                    }
                }
            }
        }
    }

    /**
     * 在JAR包中查找类
     */
    private static void findClassesInJar(URL jarUrl, String packageName, Set<Class<?>> classes) {
        String jarPath = jarUrl.getPath().substring(5, jarUrl.getPath().indexOf("!"));
        try (JarFile jarFile = new JarFile(jarPath)) {
            Enumeration<JarEntry> entries = jarFile.entries();
            while (entries.hasMoreElements()) {
                JarEntry entry = entries.nextElement();
                String entryName = entry.getName();

                if (entryName.endsWith(".class") && entryName.startsWith(packageName.replace('.', '/'))) {
                    try {
                        String className = entryName.replace('/', '.').substring(0, entryName.length() - 6);
                        classes.add(Class.forName(className));
                    } catch (ClassNotFoundException e) {
                        logger.warn("无法加载类: {}", e.getMessage());
                    }
                }
            }
        } catch (IOException e) {
            logger.error("读取JAR文件时发生错误: {}", e.getMessage(), e);
        }
    }
} 